// The main ray tracer.

#include <Fl/fl_ask.h>

#include <cmath>
#include "RayTracer.h"
#include "scene/light.h"
#include "scene/material.h"
#include "scene/ray.h"
#include "fileio/read.h"
#include "fileio/parse.h"

// Trace a top-level ray through normalized window coordinates (x,y)
// through the projection plane, and out into the scene.  All we do is
// enter the main ray-tracing method, getting things started by plugging
// in an initial ray weight of (0.0,0.0,0.0) and an initial recursion depth of 0.
vec3f RayTracer::trace( Scene *scene, double x, double y )
{
    ray r( vec3f(0,0,0), vec3f(0,0,0) );
    scene->getCamera()->rayThrough( x,y,r );
	
	return traceRay( scene, r, vec3f(1.0,1.0,1.0), traceDepth ).clamp();
}

// Do recursive ray tracing!  You'll want to insert a lot of code here
// (or places called from here) to handle reflection, refraction, etc etc.
vec3f RayTracer::traceRay( Scene *scene, const ray& r, 
	const vec3f& thresh, int depth )
{
	if( depth < 0 ) 
		return vec3f(0.0,0.0,0.0);

	isect i;

	if( scene->intersect( r, i ) ) {
		// YOUR CODE HERE

		// An intersection occured!  We've got work to do.  For now,
		// this code gets the material for the surface that was intersected,
		// and asks that material to provide a color for the ray.  

		// This is a great place to insert code for recursive ray tracing.
		// Instead of just returning the result of shade(), add some
		// more steps: add in the contributions from reflected and refracted
		// rays.

		const Material& m = i.getMaterial();
		vec3f result = m.shade(scene, r, i);

		vec3f reflectedRayResult = vec3f(0.0,0.0,0.0);
		// check if the material reflect rays
		if ( (m.kr * m.kr) != 0 ) {
			vec3f reflectDir = r.getDirection() - 2*(r.getDirection()*i.N)*i.N;
			
			ray reflectedRay = ray(r.at(i.t), reflectDir.normalize());
			reflectedRayResult = traceRay(scene, reflectedRay, thresh, depth-1);
			vec3f adjustedReflectedRayResult = vec3f(m.kr[0] * reflectedRayResult[0], 
										m.kr[1] * reflectedRayResult[1], 
										m.kr[2] * reflectedRayResult[2]);
			result += adjustedReflectedRayResult; 
		}

		// check if the material refract rays
		if ( (m.kt * m.kt) != 0 ) {
			double n; // n = n1 / n2
			double cosI;
			double normalDir;

			if ( r.getDirection() * i.N < 0 ) { // eye ray entering object, no TIR
				n = 1 / m.index;
				cosI = -1 *i.N.normalize()*r.getDirection().normalize();
				normalDir = 1;
			} else { // eye ray leaving obj, TIR may occur
				n = m.index;
				cosI = i.N.normalize()*r.getDirection().normalize();
				normalDir = -1;
			}

			//double cosI = -1 *i.N.normalize()*r.getDirection().normalize();
			double sq_sinT= n * n * (1 - cosI * cosI ) ;
				
			if (sq_sinT > 1.0 ) // TIR occurs
				return vec3f(0.0,0.0,0.0); 

			double cosT = sqrt(1 - sq_sinT);

			vec3f refractedDir = n * r.getDirection() + (n* cosI - cosT) * (normalDir * i.N);
			ray refractedRay = ray(r.at(i.t), refractedDir.normalize());
			vec3f refractedRayResult = traceRay(scene, refractedRay, thresh, depth-1);
			vec3f adjustedRefractedRayResult = vec3f( m.kt[0] * refractedRayResult[0],
										m.kt[1] * refractedRayResult[1],
										m.kt[2] * refractedRayResult[2]);

			result += adjustedRefractedRayResult;

		}

		return result;
	
	} else {
		// No intersection.  This ray travels to infinity, so we color
		// it according to the background color, which in this (simple) case
		// is just black.

		return vec3f( 0.0, 0.0, 0.0 );
	}
}

RayTracer::RayTracer()
{
	buffer = NULL;
	intensity_buffer = NULL;
	buffer_width = buffer_height = 256;
	scene = NULL;
	superSamplingSize = 1;
	traceDepth = 0;
	isJitterEnabled = false;
	isAdaptiveSamplingEnabled = false;

	m_bSceneLoaded = false;
}


RayTracer::~RayTracer()
{
	delete [] buffer;
	delete [] intensity_buffer;
	delete scene;
}

void RayTracer::getBuffer( unsigned char *&buf, int &w, int &h )
{
	buf = buffer;
	w = buffer_width;
	h = buffer_height;
}

void RayTracer::getIntensityBuffer( unsigned char *&buf, int &w, int &h )
{
	buf = intensity_buffer;
	w = buffer_width;
	h = buffer_height;
}

double RayTracer::aspectRatio()
{
	return scene ? scene->getCamera()->getAspectRatio() : 1;
}

bool RayTracer::sceneLoaded()
{
	return m_bSceneLoaded;
}

bool RayTracer::loadScene( char* fn )
{
	try
	{
		scene = readScene( fn );
	}
	catch( ParseError pe )
	{
		fl_alert( "ParseError: %s\n", pe );
		return false;
	}

	if( !scene )
		return false;
	
	buffer_width = 256;
	buffer_height = (int)(buffer_width / scene->getCamera()->getAspectRatio() + 0.5);

	bufferSize = buffer_width * buffer_height * 3;
	buffer = new unsigned char[ bufferSize ];
	intensity_buffer = new unsigned char[bufferSize];
	
	// separate objects into bounded and unbounded
	scene->initScene();
	
	// Add any specialized scene loading code here
	
	m_bSceneLoaded = true;

	return true;
}

void RayTracer::traceSetup( int w, int h )
{
	if( buffer_width != w || buffer_height != h )
	{
		buffer_width = w;
		buffer_height = h;

		bufferSize = buffer_width * buffer_height * 3;
		delete [] buffer;
		delete [] intensity_buffer;
		buffer = new unsigned char[ bufferSize ];
		intensity_buffer = new unsigned char[ bufferSize ];
	}
	memset( buffer, 0, w*h*3 );
	memset( buffer, 0, w*h*3 );
}

void RayTracer::traceLines( int start, int stop )
{
	vec3f col;
	if( !scene )
		return;

	if( stop > buffer_height )
		stop = buffer_height;

	for( int j = start; j < stop; ++j )
		for( int i = 0; i < buffer_width; ++i )
			tracePixel(i,j);
}

void RayTracer::tracePixel( int i, int j )
{
	vec3f col = vec3f(0.0,0.0,0.0);
	int numOfRays = 0;

	if( !scene )
		return;

	double subDivWidth = (superSamplingSize-1) * buffer_width;
	double subDivHeight = (superSamplingSize-1) * buffer_height;
	subDivWidth = (subDivWidth == 0) ? 1 : subDivWidth;
	subDivHeight = (subDivHeight == 0) ? 1 : subDivHeight;
	double totalDiv = superSamplingSize * superSamplingSize;
	if (isJitterEnabled) {
		for( int m = 0; m < totalDiv; m++ ) {
			double x = double(i)/double(buffer_width) + double(1.0/ ((rand() % 10 + 1) *buffer_width));
			double y = double(j)/double(buffer_height) + double(1.0/ ((rand() % 10 + 1) *buffer_height));
			col += trace( scene,x,y )/ totalDiv;
		}
	} else {
		if(isAdaptiveSamplingEnabled) {
			double threshold = 1.0/100;
			double x = double(i)/double(buffer_width);
			double y = double(j)/double(buffer_height);

			col = adaptiveTrace( scene, x, y, buffer_width, buffer_height, threshold, 0, numOfRays);


		} else {
			for (int m = 0; m < superSamplingSize; m++ ) {
				double x = double(i)/double(buffer_width) + double(m)/subDivWidth;
				for (int n = 0; n < superSamplingSize; n++ ) {
					double y = double(j)/double(buffer_height) + double(n)/subDivHeight;

					col += trace( scene,x,y )/ totalDiv;
				}

			}
		}
	}
	unsigned char *pixel = buffer + ( i + j * buffer_width ) * 3;
	if( isAdaptiveSamplingEnabled ) {
		unsigned char *intensity_map = intensity_buffer + ( i + j * buffer_width ) * 3;
		int intensity = (int) (numOfRays/64.0 * 255);
		intensity_map[0] = intensity;
		intensity_map[1] = intensity;
		intensity_map[2] = intensity;
	}


	pixel[0] = (int)( 255.0 * col[0]);
	pixel[1] = (int)( 255.0 * col[1]);
	pixel[2] = (int)( 255.0 * col[2]);


}

vec3f RayTracer::adaptiveTrace( Scene *scene, double x, double y, double divWidth, double divHeight, double threshold, int depth, int &numOfRays)
{
	numOfRays += 4;
	vec3f col1 = trace( scene, x, y );
	vec3f col2 = trace( scene, x+double(1/divWidth), y);
	vec3f col3 = trace( scene, x, y+double(1/divHeight));
	vec3f col4 = trace( scene, x+double(1/divWidth), y+double(1/divHeight));

	if( depth >= 3)
		return (col1 + col2 + col3 + col4)/4;

	if( (col1 - col2).length() < threshold && 
		(col2 - col3).length() < threshold && 
		(col3 - col4).length() < threshold && 
		(col4 - col1).length() < threshold ) {
			return col1;
	} else {
		divWidth*=2;
		divHeight*=2;
		return (adaptiveTrace(scene, x, y, divWidth, divHeight,threshold, depth+1, numOfRays) +
				adaptiveTrace(scene, x+double(1/divWidth), y, divWidth, divHeight,threshold, depth+1,numOfRays) +
				adaptiveTrace(scene, x, y+double(1/divHeight), divWidth, divHeight,threshold, depth+1,numOfRays) +
				adaptiveTrace(scene, x+double(1/divWidth), y+double(1/divHeight), divWidth, divHeight,threshold, depth+1,numOfRays) )/4;
	}

}